import CGNS.MAP as CGM
import CGNS.PAT.cgnskeywords as CGK
import CGNS.PAT.cgnsutils as CGU
import numpy as np

from service.utils.cell_type_helper import get_cell_type_by_point_count


class ICgnsReaderV4:
    reader_instance = None
    series_reader_instance = None
    attributes = None
    container = None
    cell_flag = None
    node_flag = None
    standard = None

    def __init__(self):
        self.container = {"version": "", "root": "", "zones": []}
        self.cell_flag = 0
        self.node_flag = 0
        self.standard = {
            "blocks": [],
            "nodes": [],
            "elements": [],
            "nodeIds": [],
            "attributes": {},
        }

    def read(self, path: str):
        # try:
        (tree, links, paths) = CGM.load(path)
        if len(tree) >= 3:
            info = tree[2]
            for node in info:
                self.parse_node(node)
        self.format_to_standard()

    # except Exception as e:
    #     print(f"读取 CGNS 文件时出错: {e}")

    def parse_zone_child(self, node, data):
        if node[3] == CGK.GridCoordinates_ts:
            for child_node in node[2]:
                if child_node[3] == "DataArray_t":
                    print(child_node[0])
                    if child_node[0] == "CoordinateX" or child_node[0].lower() == "x":
                        data["x"] = data["x"] + np.array(child_node[1]).tolist()
                    elif child_node[0] == "CoordinateY" or child_node[0].lower() == "y":
                        data["y"] = data["y"] + np.array(child_node[1]).tolist()
                    elif child_node[0] == "CoordinateZ" or child_node[0].lower() == "z":
                        data["z"] = data["z"] + np.array(child_node[1]).tolist()
        elif node[3] == "ZoneBC_t":
            # 边界条件不予处理
            return
        elif node[3] == "ZoneGridConnectivity_t":
            print(node[0])
        # 解析物理量
        elif node[3] == "FlowSolution_t":
            self.parse_attributes(node, data)
        # 解析单元数据
        elif node[3] == CGK.Elements_ts:
            self.parse_elements(node, data)

    def parse_elements(self, node, data):
        ele_type = ""
        ele_count = 0
        point_count = 0
        point_index_arr = []
        for child_node in node[2]:
            if child_node[0] == CGK.ElementRange_s:
                # 单元数量
                ele_count = child_node[1][1] - child_node[1][0]
            elif child_node[0] == CGK.ElementConnectivity_s:
                point_count = len(child_node[1])
                point_index_arr = child_node[1].copy()

        per_point = int(point_count / ele_count)
        ele_type = get_cell_type_by_point_count(per_point, node[0])
        for i in range(ele_count):
            node_ids = []
            for j in range(per_point):
                # -1 是因为 CGNS 文件中的索引从 1 开始，而 Python 中的索引从 0 开始
                node_ids.append(int(point_index_arr[i * per_point + j] - 1))
            ele_info = {
                "type": ele_type,
                "eid": i,
                "nodeIds": node_ids,
                "faceIds": [i],
                "faces": [],
                "physicalTag": 0,
            }
            match ele_type:
                case "CHEXA":
                    ele_info["faces"] = [
                        [0, 1, 2, 3],
                        [4, 5, 6, 7],
                        [0, 1, 5, 4],
                        [1, 2, 6, 5],
                        [2, 3, 7, 6],
                        [3, 0, 4, 7],
                    ]
                    ele_info["dim"] = 3
                case "CPENTA":
                    ele_info["faces"] = [
                        [0, 1, 2],
                        [3, 4, 5],
                        [0, 1, 4, 3],
                        [1, 2, 5, 4],
                        [2, 0, 3, 5],
                    ]
                    ele_info["dim"] = 3

            data["elements"].append(ele_info)

        print(f"每个单元 {per_point} 个点", ele_type, node[0])

    def parse_node(self, node):
        if node[3] == "CGNSLibraryVersion_t":
            self.container["version"] = node[2]
        elif node[3] == "CGNSBase_t":
            self.container["root"] = node[0]
            for child in node[2]:
                self.parse_node(child)
        elif node[3] == "Zone_t":
            t = {
                "name": node[0],
                "elements": [],
                "attributes": {},
                "x": [],
                "y": [],
                "z": [],
            }
            self.container["zones"].append(t)
            self.standard["blocks"].append(
                {
                    "name": node[0],
                    "elements": [],
                }
            )
            for child in node[2]:
                self.parse_zone_child(
                    child,
                    t,
                )
        elif node[3] == "Family_t":
            for child in node[2]:
                self.parse_node(child)
        elif node[3] == "ReferenceState_t":
            print(node)

    def parse_attributes(self, node, data):
        for child_node in node[2]:
            data["attributes"][child_node[0]] = child_node[1]

    def get_version(self, path):
        try:
            (tree, links, paths) = CGM.load(path)
            if len(tree) >= 3:
                info = tree[2]
                for node in info:
                    if node[3] == "CGNSLibraryVersion_t":
                        self.container["version"] = node[1]
            return self.container["version"]
        except Exception as e:
            return "3"

    # 转换成标准数据格式
    def format_to_standard(self):
        z_index = 0
        for zone in self.container["zones"]:
            for i in range(len(zone["x"])):
                self.standard["nodes"].append(zone["x"][i])
                self.standard["nodes"].append(zone["y"][i])
                self.standard["nodes"].append(zone["z"][i] if len(zone["z"]) > 0 else 0)
                self.standard["nodeIds"].append(i + self.node_flag)
            self.node_flag += len(zone["x"])
            for ele in zone["elements"]:
                ele["eid"] = ele["eid"] + self.cell_flag
                ele["physicalTag"] = z_index
                self.standard["elements"].append(ele)
            # 注入范围
            self.standard["blocks"][z_index]["bid"] = z_index
            self.standard["blocks"][z_index]["elements"].append(self.cell_flag)
            self.cell_flag += len(zone["elements"])
            self.standard["blocks"][z_index]["elements"].append(self.cell_flag)
            attrs: dict = self.standard["attributes"]
            for key in zone["attributes"].keys():
                if attrs.get(key) is None:
                    attrs[key] = {
                        "name": key,
                        "size": 1,
                        "dx": [],
                        "dy": [],
                        "dz": [],
                        "magtitude": [],
                    }
                attrs[key]["dx"] = attrs[key]["dx"] + (
                    np.array(zone["attributes"][key]).tolist()
                )
            z_index += 1

    def get_attribute_list(self):
        return self.standard["attributes"].keys()
